import paddle.fluid.layers as L

from helm.static.models.layers import Conv2d, GlobalAvgPool, Pool2d, Identity, Linear, Layer, Sequential
from helm.static.models.modules import SELayer

__all__ = ["RegNet"]


class Bottleneck(Layer):

    def __init__(self, in_channels, out_channels, stride, groups, use_se):
        super().__init__()
        self.use_se = use_se

        self.conv1 = Conv2d(in_channels, out_channels, kernel_size=1, bn=True, act='default')
        self.conv2 = Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, groups=groups,
                            bn=True, act='default')
        if self.use_se:
            self.se = SELayer(out_channels, 4)
        self.conv3 = Conv2d(out_channels, out_channels, kernel_size=1, bn=True, zero_init=True)
        if stride != 1 or in_channels != out_channels:
            layers = []
            if stride != 1:
                layers.append(Pool2d(kernel_size=(2, 2), stride=2, type='avg'))
            layers.extend([
                Conv2d(in_channels, out_channels, kernel_size=1, bn=True),
            ])
            self.shortcut = Sequential(*layers)
        else:
            self.shortcut = Identity()

    def forward(self, x):
        identity = self.shortcut(x)
        x = self.conv1(x)
        x = self.conv2(x)
        if self.use_se:
            x = self.se(x)
        x = self.conv3(x)
        x = L.elementwise_add(x, identity, act='default')
        return x


def _make_layer(in_channels, out_channels, num_units, stride, groups, use_se):
    layers = [("unit1", Bottleneck(in_channels, out_channels, stride=stride, groups=groups, use_se=use_se))]
    for i in range(1, num_units):
        layers.append(
            ("unit%d" % (i + 1), Bottleneck(out_channels, out_channels, stride=1, groups=groups, use_se=use_se)))
    return Sequential(*layers)


class RegNet(Layer):

    def __init__(self, stem_channels=32, channels_per_stage=(96, 256, 640), units_per_stage=(5, 5, 5),
                 channels_per_group=16, use_se=True, num_classes=10):
        super().__init__()

        self.conv = Conv2d(3, stem_channels, kernel_size=3, bn=True, act='default')

        cs = channels_per_stage
        gs = [c // channels_per_group for c in cs]
        us = units_per_stage

        self.stage1 = _make_layer(
            stem_channels, cs[0], us[0], stride=1, groups=gs[0], use_se=use_se)
        self.stage2 = _make_layer(
            cs[0], cs[1], us[1], stride=2, groups=gs[1], use_se=use_se)
        self.stage3 = _make_layer(
            cs[1], cs[2], us[2], stride=2, groups=gs[2], use_se=use_se)

        self.avgpool = GlobalAvgPool()
        self.fc = Linear(cs[2], num_classes)

    def forward(self, x):
        x = self.conv(x)

        x = self.stage1(x)
        x = self.stage2(x)
        x = self.stage3(x)

        x = self.avgpool(x)
        x = self.fc(x)
        return x


# def test_regnet():
#     x = torch.randn(1, 3, 32, 32)
#
#     net = RegNet(32, (96, 256, 640), [4, 8, 2], 16, use_se=True)
#     assert profile(net, (x,))[1] == 3920266
